import pandas as pd
import statsmodels.api as sm
import os
import numpy as np

import matplotlib.pyplot as plt

from utils.data_split import DataSplitUtil
from utils.utils import *  # 包含 target_column、exclude_columns、regression_result_path、opj、md
import utils.data_split
# 读取数据
df = pd.read_csv(signature_score_csv_path)
dp = DataSplitUtil(split_random_state=split_random_state_list[0])
train,test = dp.get_train_test_df(df)
Y = train[target_column]
X_all = train.drop(columns=exclude_columns)

# 创建结果目录
md(regression_result_path)

# ------------------------
# 🔹 单变量逻辑回归分析
# ------------------------

results_list = []

for col in X_all.columns:
    X_single = sm.add_constant(X_all[[col]])
    try:
        model_single = sm.Logit(Y, X_single).fit(disp=0)
        or_value = np.exp(model_single.params[col])
        p_val = model_single.pvalues[col]
        star = '***' if p_val < 0.001 else '**' if p_val < 0.01 else '*' if p_val < 0.05 else ''
        results_list.append({
            'Feature': col,
            'Pseudo_R_squared': model_single.prsquared,
            'OR': or_value,
            'P_value': p_val,
            'P_star': f"{p_val:.3e}{star}"
        })
    except Exception as e:
        print(f"❌ 特征 {col} 回归失败：{e}")

# 保存单变量逻辑回归结果
single_df = pd.DataFrame(results_list)
single_save_path = opj(regression_result_path,'single')
md(single_save_path)
single_df.to_csv(opj(single_save_path,"logistic_regression_results.csv"), index=False)

# ------------------------
# 🔹 图1: Pseudo R² 柱状图
# ------------------------
plt.figure(figsize=(8, 5))
pseudo_r2_sorted = single_df.sort_values("Pseudo_R_squared", ascending=False)
plt.barh(pseudo_r2_sorted["Feature"], pseudo_r2_sorted["Pseudo_R_squared"], color="skyblue")
plt.xlabel("Pseudo R²")
plt.title("Pseudo R² for Single-variable Logistic Regression")
plt.gca().invert_yaxis()
plt.tight_layout()
plt.savefig(opj(single_save_path, "pseudo_r2_barplot.pdf"))
plt.close()

# ------------------------
# 🔹 图2: OR 条形图 + 星号标注
# ------------------------
plt.figure(figsize=(8, 5))
or_sorted = single_df.sort_values("OR", ascending=False)
bars = plt.barh(or_sorted["Feature"], or_sorted["OR"], color="salmon")
for i, (val, star) in enumerate(zip(or_sorted["OR"], or_sorted["P_star"])):
    plt.text(val, i, f' {star}', va='center')
plt.xlabel("Odds Ratio (OR)")
plt.title("Odds Ratios with Significance")
plt.gca().invert_yaxis()
plt.tight_layout()
plt.savefig(opj(single_save_path, "or_barplot_with_pstars.pdf"))
plt.close()

single_df.head()


multi_save_path = opj(regression_result_path,'multi')
md(multi_save_path)
# 多变量逻辑回归分析，使用 add_constant 并拟合
X_multi = sm.add_constant(X_all)
model_multi = sm.Logit(Y, X_multi).fit(disp=0)

# 提取 OR、P 值、伪R²
multi_df = pd.DataFrame({
    "Feature": model_multi.params.index,
    "Coef": model_multi.params.values,
    "P_value": model_multi.pvalues.values
})

# 去掉常数项
multi_df = multi_df[multi_df["Feature"] != "const"].copy()

# 添加 OR 和 星号标注
multi_df["OR"] = np.exp(multi_df["Coef"])
multi_df["P_star"] = multi_df["P_value"].apply(
    lambda p: f"{p:.3e}" + ("***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else "")
)
multi_df["Pseudo_R_squared"] = model_multi.prsquared

# 保存结果
multi_df = multi_df[["Feature", "Pseudo_R_squared", "OR", "P_value", "P_star"]]
multi_path = opj(multi_save_path, "multivariable_logistic_regression_results_OR.csv")
multi_df.to_csv(multi_path, index=False)

# 绘图: OR 条形图 + 星号
plt.figure(figsize=(8, 5))
multi_sorted = multi_df.sort_values("OR", ascending=False)
plt.barh(multi_sorted["Feature"], multi_sorted["OR"], color="mediumseagreen")
for i, (val, star) in enumerate(zip(multi_sorted["OR"], multi_sorted["P_star"])):
    plt.text(val, i, f' {star}', va='center')
plt.xlabel("Odds Ratio (OR)")
plt.title("Multivariable Logistic Regression ORs with Significance")
plt.gca().invert_yaxis()
plt.tight_layout()
plt.savefig(opj(multi_save_path, "multi_or_barplot_with_pstars.pdf"))
plt.close()

multi_df.head()
